import deepxde as dde
import numpy as np
from matplotlib import pyplot as plt
import os

main_path = os.getcwd()

# 需要对哪个level的网格进行加密？
level = 2

"""------------------------------ Fluid parameters  -----------------------"""
# The universal gas constant 气体常数 (FOR AIR HERE)
R = 286.9 #(J/Kg.K)

# 空气绝热指数
k = 1.4

# 用于无量纲化的速度、压力和密度，取入口处的声速值、压力值和密度值
p0 = 73048 #(pa)
rho0 = 0.9 #(kg/m^3)
a0 = np.sqrt(k*p0/rho0) #(m/s)

"""------------------------------ NACA001065 shape -----------------------"""
airfoil_boundary_points = np.load("../0_Mesh_files/ordered_airfoil_points.npy")[:, 0:2]
airfoil = dde.geometry.Polygon(airfoil_boundary_points)

"""------------------------------------ Training Spatio-temporal region --------------------------------------"""
"""-------------- Rectangular -------------"""
x_min, x_max = -5.5, 6.5
L = x_max - x_min

y_min, y_max = -2.0, 2.0
H = y_max - y_min

"""Spatial region definition"""
farfield = dde.geometry.Rectangle([x_min, y_min], [x_max, y_max])  # 设置左下角和右上角坐标
geom = dde.geometry.CSGDifference(farfield, airfoil)

"""------------------------------ Dimensionless Euler Equations(Time-dependent PDEs) -----------------------"""
def DL_Euler_Equation_2D(x, y):
    """
    System of PDEs to be minimized: compressible 2D Euler equations.

    """
    gamma = 1.4
    
    u = y[:, 0:1]
    v = y[:, 1:2]
    p = y[:, 2:3]
    rho = y[:, 3:4]

    E = p/(gamma-1) + 0.5*rho*(u*u + v*v)

    A1, B1 = rho*u, rho*v
    A2, B2 = gamma*rho*u*u + p, gamma*rho*u*v
    A3, B3 = gamma*rho*u*v, gamma*rho*v*v + p
    A4, B4 = u*E, v*E

    dA1_x = dde.grad.jacobian(A1, x, i=0, j=0)
    dB1_y = dde.grad.jacobian(B1, x, i=0, j=1)

    dA2_x = dde.grad.jacobian(A2, x, i=0, j=0)
    dB2_y = dde.grad.jacobian(B2, x, i=0, j=1)

    dA3_x = dde.grad.jacobian(A3, x, i=0, j=0)
    dB3_y = dde.grad.jacobian(B3, x, i=0, j=1)

    dA4_x = dde.grad.jacobian(A4, x, i=0, j=0)
    dB4_y = dde.grad.jacobian(B4, x, i=0, j=1)
    
    continuity = dA1_x + dB1_y
    x_momentum = dA2_x + dB2_y
    y_momentum = dA3_x + dB3_y
    energy = dA4_x + dB4_y
    return [continuity, x_momentum, y_momentum, energy]

"""---------------------------------------- Boundary Definition -------------------------------------------"""
"""Boundary definition of rectangle domain"""
def boundary_inlet(x, on_boundary):
    """Checks the inlet boundary"""
    return np.isclose(x[0], x_min)

def boundary_outlet(x, on_boundary):
    """Checks the outlet boundary"""
    return np.isclose(x[0], x_max)

def boundary_top_bottom(x, on_boundary):
    """Checks the bottom boundary"""
    return (np.isclose(x[1], y_min) or np.isclose(x[1], y_max))

"""Boundary definition of cylinder domain"""
def boundary_airfoil(x, on_boundary):
    return on_boundary and (not farfield.on_boundary(x))

'''-----------------------Defining BC for u, v, and p (component = 0, 1, 2 resp assign)---------------------------
------------------------------
Dirichlet boundary conditions:
    y(x) = func(x)
------------------------------
Neumann boundary conditions: 
    dy/dn(x) = func(x)
------------------------------'''
# Boundary values definition
def No_slip(x): # 无滑移
    return 0.0

def uniform_flow(x): # 无滑移
    return 0.7

def Zero_v(x):  # 零速度
    return 0.0

def Zero_p(x):  # 零压力
    return 0.0

'''----------------------- Other functions ---------------------------'''
def gen_training_points():
    mesh_xy_points = data_points[:, 0:2]
    return mesh_xy_points

def plot_points_xy(points, color="k", marker="."):
    figure = plt.figure(dpi=300)
    axis = figure.add_subplot(111)
    axis.scatter(points[:, 0], points[:, 1], s=0.5, color=color, marker=marker)
    plt.axis("scaled")  # 设置x轴和y轴相同的缩放比例
    plt.savefig('./Figs/Total_points.png', dpi=300)
    plt.show()

def plot_flow_field():
    import matplotlib as mpl
    from matplotlib import rcParams
    config = {"font.size": 16}
    rcParams.update(config)
    plt.rcParams['xtick.direction'] = 'in'
    plt.rcParams['ytick.direction'] = 'in'
    '''------------------------------------------- Model validation ---------------------------------------------'''
    '''------------------Plotting 2------------------'''

    test_data = data_xy

    # Model predictions generation
    u = model.predict(test_data)[:, 0] * a0
    v = model.predict(test_data)[:, 1] * a0
    p = model.predict(test_data)[:, 2] * p0
    rho = model.predict(test_data)[:, 3] * rho0

    airfoil_plot = airfoil_boundary_points[:, 0:2]

    # plot
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], (u**2 + v**2)**0.5, levels=1000, cmap="coolwarm", vmin=48.9, vmax=462.5)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=48.9, vmax=462.5)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Velocity Magnitude')
    plt.savefig("./Figs/Velocity.png", dpi=300)
    plt.show()   
    
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], p, levels=1000, cmap="coolwarm", vmin=19587.33, vmax=98563.27)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=19587.33, vmax=98563.27)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[Pa]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Pressure')
    plt.savefig("./Figs/P.png", dpi=300)
    plt.show()  
    
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.tricontourf(data_xy[:, 0], data_xy[:, 1], rho, levels=1000, cmap="coolwarm", vmin=0.315, vmax=1.109)
    plt.axis('scaled')
    plt.fill(airfoil_plot[:, 0], airfoil_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=0.315, vmax=1.109)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap="coolwarm"), ax=ax, label='[kg/m^3]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Density')
    plt.savefig("./Figs/Rho.png", dpi=300)
    plt.show()  
     

def make_dir():
    # 保存训练结果需要的文件
    # 模型保存地址
    Model_path = "Model/"  
    # 图片保存地址
    Figs_path = "Figs/"  
    # Loss保存地址
    Train_history_path = "Train_history/"  

    if not os.path.exists(Model_path):
    	os.mkdir(Model_path)
    if not os.path.exists(Figs_path):
    	os.mkdir(Figs_path)
    if not os.path.exists(Train_history_path):
    	os.mkdir(Train_history_path)

def plot_loss_history(loss_history, fname=None):
    loss_train = np.sum(loss_history.loss_train, axis=1)
    loss_test = np.sum(loss_history.loss_test, axis=1)

    plt.figure()
    plt.semilogy(loss_history.steps, loss_train, label="Train loss")
    plt.semilogy(loss_history.steps, loss_test, label="Test loss")
    for i in range(len(loss_history.metrics_test[0])):
        plt.semilogy(
            loss_history.steps,
            np.array(loss_history.metrics_test)[:, i],
            label="Test metric",
        )
    plt.xlabel("# Steps")
    plt.legend()

    if isinstance(fname, str):
        plt.savefig(fname, dpi=300)



'''-------------------------------------------------------------- training model:  ------------------------------------------------------------------'''
dde.config.set_random_seed(1234)
dde.config.set_default_float("float32")

make_dir()

'''------------------BC------------------'''
# data
data_points = np.load('../1_Fluent_results/level' + str(level) + '_fluent_results.npy')

# 此处注意无量纲化
data_xy = data_points[:, 0:2]
data_u = data_points[:, 2:3]/a0
data_v = data_points[:, 3:4]/a0
data_p = data_points[:, 4:5]/p0
data_rho = data_points[:, 5:6]/rho0

bc_data_u = dde.icbc.PointSetBC(data_xy, data_u, component = 0)
bc_data_v = dde.icbc.PointSetBC(data_xy, data_v, component = 1)
bc_data_p = dde.icbc.PointSetBC(data_xy, data_p, component = 2)
bc_data_rho = dde.icbc.PointSetBC(data_xy, data_rho, component = 3)

bcs = [
      bc_data_u,
      bc_data_v,
      bc_data_p,
      bc_data_rho,
]
'''--------------- Training Model ----------------'''
# 根据目前的geomtime生成训练点
total_points = gen_training_points()
plot_points_xy(total_points)

'''----Training datasets and Loss function ----'''
data = dde.data.PDE(geom,
              DL_Euler_Equation_2D, # Loss function
              bcs,                 # IC and BC
              num_domain = 0,      # training data
              num_boundary = 0,     # boundary data
              num_test = 1000,         # test data
              anchors = total_points,
              )         
print(data.num_bcs)
print(data.train_x_all.shape[0])

'''--------- Neural Network setup ---------'''
layer_size = [2] + [40] * 6 + [4]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)

'''-------- Compile, Training and save Model --------'''
model = dde.Model(data, net)
weights = [1] * 4 + [10] * 4

# first use Adam to find a roughly good solution, 
model.compile("adam", lr=1e-3, loss_weights = weights)
loss_history, train_state = model.train(epochs=20000, display_every=1000)
plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
model.save(save_path="./Model/model")
plot_flow_field()

model.compile("adam", lr=5e-4, loss_weights = weights)
loss_history, train_state = model.train(epochs=20000, display_every=1000)
plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
model.save(save_path="./Model/model")
plot_flow_field()

model.compile("adam", lr=1e-4, loss_weights = weights)
loss_history, train_state = model.train(epochs=20000, display_every=1000)
plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
model.save(save_path="./Model/model")
plot_flow_field()

# and then use L-BFGS to search a local minimum.
model.compile(optimizer="L-BFGS", loss_weights = weights)
model.train_step.optimizer_kwargs = {
    "options": {
        "maxcor": 50,
        "ftol": 1.0 * np.finfo(float).eps,
        "maxfun": 50000,
        "maxiter": 50000,
        "maxls": 50,
    }
}
# Optional: Save the model during training.
checkpointer = dde.callbacks.ModelCheckpoint("Model/model", verbose=1000, save_better_only=True)
loss_history, train_state = model.train(display_every=100, callbacks=[checkpointer])
plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
model.save(save_path="./Model/model")
plot_flow_field()



